package com.sofun.ai.controller;

import com.sofun.ai.assistant.RagAssistant;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;


/**
 * RagAssistant配置了embedding，用户输入的问题都是通过embedding的数据来回答
 * */

@RestController()
@RequestMapping("/rag")
public class RagChatController {

    @Resource
    QwenChatModel qwenChatModel;

    @Resource
    RagAssistant ragAssistant;

    @RequestMapping(value = "/memory_chat")
    public String memoryChat(@RequestParam(defaultValue="我叫张佳伟") String message) {
        return ragAssistant.chat(message);
    }

    /**
     * @param webSearchEnable 是否使用web联网搜索
     * */
    @RequestMapping(value = "/memory_stream_chat",produces ="text/stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam(defaultValue="我是谁") String message,
                                         @RequestParam(value = "webSearchEnable", required = false, defaultValue = "false") Boolean webSearchEnable) {
        TokenStream stream = ragAssistant.stream(message, LocalDate.now().toString());

        return Flux.create(sink -> stream.onPartialResponse(sink::next)
                .onCompleteResponse(c -> sink.complete())
                .onError(sink::error)
                .start());
    }

    /**
     * 流式聊天
     *
     * @param memoryId       会话id
     * @param role            设定角色
     * @param question        原始问题
     * @param webSearchEnable 是否开启网页搜索
     * @param extraInfo       额外信息（暂未实现）
     * @return
     */
    @GetMapping(value = "/stream/flux", produces ="text/stream;charset=UTF-8")
    public Flux<String> chatStreamFlux(@RequestParam(value = "memoryId") String memoryId,
                                       @RequestParam(value = "role", required = false, defaultValue = "智能问答助手") String role,
                                       @RequestParam(value = "question") String question,
                                       @RequestParam(value = "webSearchEnable", required = false, defaultValue = "false") Boolean webSearchEnable,
                                       @RequestParam(value = "extraInfo", required = false, defaultValue = "") String extraInfo) {

        TokenStream stream = ragAssistant.chatStreamTokenStream(memoryId, role, question, extraInfo);

        return Flux.create(sink -> stream.onPartialResponse(sink::next)
                .onCompleteResponse(c -> sink.complete())
                .onError(sink::error)
                .start());
    }

}
